package com.mxd.flink.connector.table;

import com.mxd.flink.connector.config.Delimiters;
import com.mxd.flink.connector.config.RedisCmd;
import com.mxd.flink.connector.config.RedisOptions;
import com.mxd.flink.connector.operator.RedisOperator;
import com.mxd.flink.connector.operator.RedisOperators;
import com.mxd.flink.connector.utils.RedisRowUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.calcite.shaded.com.google.common.cache.Cache;
import org.apache.flink.calcite.shaded.com.google.common.cache.CacheBuilder;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.functions.AsyncTableFunction;
import org.apache.flink.table.functions.FunctionContext;
import org.apache.flink.table.functions.TableFunction;
import org.apache.flink.table.types.DataType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * 异步的Loopup算子
 * @author rongdi
 * @date 2022/10/12 21:27
 */
public class AsyncRedisLookupFunction extends AsyncTableFunction<RowData> {

    private static final Logger LOG = LoggerFactory.getLogger(AsyncRedisLookupFunction.class);

    /**
     * 缓存对象
     */
    private Cache<String, Object> cache;

    /**
     * flinksql with里的参数列表
     */
    private ReadableConfig config;

    /**
     * flink中创建的redis表的数据类型集合
     */
    private List<DataType> dataTypes;

    /**
     * redis操作对象
     */
    private RedisOperator redisOperator;

    /**
     * 缓存最大条数
     */
    private long cacheMaxSize;

    /**
     * 缓存超时时间，单位秒
     */
    private long cacheTtl;

    /**
     * hget命令是否加载所有
     */
    private boolean loadAll;

    /**
     * 失败最大重试次数
     */
    private int maxRetries;

    /**
     * 访问的redis的命令
     */
    private final String command;

    private final DeserializationSchema<RowData> deserializer;


    public AsyncRedisLookupFunction(ReadableConfig config, ResolvedSchema resolvedSchema, DeserializationSchema<RowData> deserializer) {
        this.config = config;
        /**
         * 获取当前维表的表结构，其实就是一个ArrayList<TableColumn$PhysicalColumn>
         */
        this.dataTypes = resolvedSchema.getColumnDataTypes();
        this.command = config.get(RedisOptions.COMMAND);
        this.cacheMaxSize = config.get(RedisOptions.LOOKUP_CACHE_MAX_ROWS);
        this.cacheTtl = config.get(RedisOptions.LOOKUP_CHCHE_TTL);
        this.loadAll = config.get(RedisOptions.LOOKUP_CACHE_LOAD_ALL);
        this.maxRetries = config.get(RedisOptions.LOOKUP_MAX_RETRIES);
        this.deserializer = deserializer;
    }

    @Override
    public void open(FunctionContext context) throws Exception {
        /**
         * 初始化缓存，如果是cacheMaxSize和cacheTtl是默认值，表示不使用缓存
         */
        this.cache = ((this.cacheMaxSize == -1L) || (this.cacheTtl == -1L) ? null : CacheBuilder.newBuilder()
                .expireAfterWrite(this.cacheTtl, TimeUnit.SECONDS)
                .maximumSize(this.cacheMaxSize).build());
        this.redisOperator = RedisOperators.get(config);
        LOG.info("AsyncRedisLookupFunction Open");
    }

    /**
     * 核心方法，这里的keys是一个可变参数，key传过来的顺序与表申明的字段顺序有关
     * @param keys
     */
    public void eval(CompletableFuture<Collection<RowData>> resultFuture, Object... keys) throws Exception {

        /**
         * 如果缓存对象不为空，表示使用缓存
         */
        if (this.cache != null) {
            /**
             * 这里就是表示一行数据，里面就是一个Object[] fields
             */
            GenericRowData genericRowData = null;
            /**
             * 这里command是get，所以只会有一个key传进来，就是join on后面唯一的那个条件的值
             * 流表通过ScanFunction每一行数据都会调用eval方法传入key获取到value值，这里使用
             * 一个guava的缓存，先从缓存中获取维表中一行数据
             */
            if(RedisCmd.GET.equalsIgnoreCase(command)) {
                genericRowData = (GenericRowData)this.cache.getIfPresent(String.valueOf(keys[0]));
            } else if(RedisCmd.HGET.equalsIgnoreCase(command)) {
                /**
                 * 这里command是hget，原本需要两个参数key，field一起才能获取唯一值，但是为了不要每一条数据都请求一次，
                 * 这里提供一个配置项，是否根据key加载所有的field的值，同样先从缓存中获取，获取到key里存放的值后，组装
                 * 成GenericRowData对象用collect吐出去
                 */
                if (this.loadAll) {
                    Map<String, String> map = (Map) this.cache.getIfPresent(String.valueOf(keys[0]));
                    if (map != null) {
                        createRowData(resultFuture, keys, map);
                        return;
                    }
                } else {
                    /**
                     * 如果不加载key里的所有值，则直接从缓存中拿组装好的一行genericRowData
                     */
                    String key = keys[0] + Delimiters.UNDER_LINE + keys[1];
                    genericRowData = (GenericRowData)this.cache.getIfPresent(key);
                }
            }
            /**
             * 如果能从缓存拿到这一行数据，则使用complete吐出去
             */
            if (genericRowData != null) {
                resultFuture.complete(Arrays.asList(genericRowData));
                return;
            }

        }

        /**
         * 到了这里，说明缓存中没拿到，这里根据配置项的最大重试次数做个循环，然后真正去redis中取值
         */
        for (int i = 0; i <= this.maxRetries; i++) {
            try {
                query(resultFuture,keys);
                return;
            } catch (Exception e) {
                LOG.error("query redis error, retry times:{}", Integer.valueOf(i), e);
                if (i >= this.maxRetries) {
                    throw new RuntimeException("query redis error ", e);
                }
                try {
                    TimeUnit.MILLISECONDS.sleep(500 * i);
                } catch (InterruptedException e1) {
                }
            }
        }

    }

    @Override
    public void close() throws Exception {
        if (this.redisOperator != null) {
            this.redisOperator.close();
        }
        LOG.info("AsyncRedisLookupFunction Close");
    }

    private void query(CompletableFuture<Collection<RowData>> resultFuture,Object... keys) throws Exception {

        if(RedisCmd.GET.equalsIgnoreCase(command)) {
            /**
             * 这里command是get，所以传过来的keys肯定只有一个key，相当于redis里执行命令 get key
             * 很明显如果有数据，就可以获取到一个值，这里返回一个future对象，在回调中处理拿到数据之后的逻辑
             */
            CompletableFuture future = this.redisOperator.asyncGet(String.valueOf(keys[0]));
            future.whenCompleteAsync((data,exp) -> {
                try {
                    GenericRowData rowData = null;
                    /**
                     * 如果反序列化工具为空，意味着format配置项为空，则手动构造RowData
                     */
                    if (this.deserializer == null) {
                        /**
                         * 这里构造一个GenericRowData对象，由于这种情况这一行只有key、value所以只需要构造一个
                         * 可以放2个数据的对象，第一数据就是传进来的key，第二个数据就是查出来的value。这里根据上面
                         * 获取到的维表的表结构对象dataTypes拿到value的数据类型将从redis中获取到的值转换成对象
                         */
                        rowData = new GenericRowData(2);
                        rowData.setField(0, keys[0]);
                        rowData.setField(1, RedisRowUtil.dataTypeFromString(
                                this.dataTypes.get(1).getLogicalType(), String.valueOf(data)));
                    } else {
                        /**
                         * 如果反序列化工具不为空，则format配置项有内容，则直接使用反序列化工具反序列化成GenericRowData
                         * 行数据，这里需要注意的是json的内容一定要匹配上申明表的结构，也就是说json里最好要包含表结构申明的所有
                         * 列才行，当然要包含作为维表查询的join on后面带的作为key的查询列了，不然会报空指针异常
                         */
                        rowData = (GenericRowData) deserializer.deserialize(String.valueOf(data).getBytes());
                    }
                    /**
                     * 吐出封装好的GenericRowData对象，并放入缓存，这里也是上面获取缓存的时候直接取
                     * GenericRowData对象的原因
                     */
                    resultFuture.complete(Arrays.asList(rowData));
                    if ((this.cache != null) && (data != null)) {
                        this.cache.put(String.valueOf(keys[0]), rowData);
                    }
                } catch (Exception e) {
                    LOG.error("handle async get error",e);
                }
            });

        } else if(RedisCmd.HGET.equalsIgnoreCase(command)) {
            /**
             * 这里command是hget，这里同样准寻配置项，是否加载key下的所有field值，相当于执行redis命令
             * hgetall key
             */
            if (this.loadAll) {
                hgetAll(resultFuture, keys);
                return;
            }
            /**
             * 这里command是hget，所以传过来的keys肯定有两个，key和field，肯定是维表join的时候on后面
             * 带了两个条件，相当于redis里执行命令 hget key field，很明显如果有数据，就可以获取到一个值，
             * 这里返回一个future对象，在回调中处理拿到返回数据后的逻辑
             */
            CompletableFuture<String> future = this.redisOperator.asyncHget(
                    String.valueOf(keys[0]), String.valueOf(keys[1]));
            future.whenCompleteAsync((data,exp) -> {
                try {
                    GenericRowData rowData = null;
                    /**
                     * 如果反序列化工具为空，意味着format配置项为空，则手动构造RowData
                     */
                    if (this.deserializer == null) {
                        /**
                         * 这里由于keys里有两个值key和field，再加上一个value就有三个值了，所以需要构造一个可以放3个
                         * field值的GenericRowData对象，同样第三个元素value需要根据属性类型转换一下，因为申明表结构
                         * 时肯定不止string，有可能有int，long这些类型
                         */
                        rowData = new GenericRowData(3);
                        rowData.setField(0, keys[0]);
                        rowData.setField(1, keys[1]);
                        rowData.setField(2, RedisRowUtil.dataTypeFromString(
                                (this.dataTypes.get(2)).getLogicalType(), String.valueOf(data)));
                    } else {
                        /**
                         * 如果反序列化工具不为空，则format配置项有内容，则直接使用反序列化工具反序列化成GenericRowData
                         * 行数据，这里需要注意的是json的内容一定要匹配上申明表的结构，也就是说json里要包含作为维表查询的
                         * join on后面带的作为key的查询列，不然会报空指针异常
                         */
                        rowData = (GenericRowData) deserializer.deserialize(String.valueOf(data).getBytes());
                    }
                    /**
                     * 将封装好的GenericRowData对象通过complete吐出去，然后拼接key和field作为key值将该对象
                     * 放入本地缓存中，这也是为啥上面直接从缓存中可以拿到GenericRowData对象的原因
                     */
                    resultFuture.complete(Arrays.asList(rowData));
                    if ((this.cache != null) && (data != null)) {
                        String key = keys[0] + Delimiters.UNDER_LINE + keys[1];
                        this.cache.put(key, rowData);
                    }
                } catch (Exception e) {
                    LOG.error("handle async hget error",e);
                }
            });

        }
    }

    /**
     * 执行hgetall key命令获取key下的所有field值，将获取到的map放入缓存，这也是为啥
     * 上面通过缓存获取到的是一个map的原因
     * @param keys
     */
    void hgetAll(CompletableFuture<Collection<RowData>> resultFuture, Object... keys) throws Exception{
        CompletableFuture<Map<String, String>> future = this.redisOperator.asyncHgetAll(String.valueOf(keys[0]));
        future.whenCompleteAsync((data,exp) -> {
            try {
                if (data == null) {
                    return;
                }
                this.cache.put(String.valueOf(keys[0]), data);
                createRowData(resultFuture, keys, data);
            } catch (Exception e) {
                LOG.error("handle async hgetAll error",e);
            }
        });

    }

    /**
     * 这里时创建一个GenericRowData对象，并通过complete方法吐出去
     * @param keys
     * @param map
     */
    private void createRowData(CompletableFuture<Collection<RowData>> resultFuture, Object[] keys, Map<String, String> map) throws Exception{
        String value= map.get(String.valueOf(keys[1]));
        if(this.deserializer == null) {
            GenericRowData genericRowData = new GenericRowData(3);
            genericRowData.setField(0, keys[0]);
            genericRowData.setField(1, keys[1]);
            genericRowData.setField(2, RedisRowUtil.dataTypeFromString(
                    this.dataTypes.get(2).getLogicalType(), value));
            resultFuture.complete(Arrays.asList(genericRowData));
        } else {
            if(StringUtils.isEmpty(value)) {
                return;
            }
            /**
             * 拿到对应的值，然后使用反序列化工具转成GenericRowData，使用complete吐出去
             */
            GenericRowData genericRowData = (GenericRowData)this.deserializer.deserialize(value.getBytes());
            resultFuture.complete(Arrays.asList(genericRowData));
        }

    }

}
